from PIL import Image, ImageDraw, ImageColor
import fitz
from io import BytesIO
from os.path import dirname, abspath, basename
from shutil import copyfile
from tempfile import mkstemp

def colorimage(*, pdf_path, highlight, color = "#ffff007f", **args):
    ret = ""
    _nf, nfm = mkstemp(suffix=basename(pdf_path), dir=dirname(pdf_path))
    # copyfile(pdf_path, nfm)
    pd = fitz.Document(pdf_path)
    # Loop through the highlights
    for highlight in highlight:
        # Get the page object using the xref value
        pix = fitz.Pixmap(pd, highlight['xref'])
        img = Image.open(BytesIO(pix.tobytes()))
        page = ImageDraw.Draw(img, "RGBA")
        # Loop through the rectangles
        for rect in highlight['rects']:
            # Create a rectangle object
            rect = {k:float(rect[k]) for k in rect}
            r = (rect['x'], rect['y'], rect['x'] + rect['w'], rect['y'] + rect['h'])
            # Draw the rectangle on the page
            page.rectangle(r, fill=ImageColor.getrgb(color))
        _imgfp, imgtmp = mkstemp(f'{highlight["xref"]}.png',None,dirname(pdf_path))
        img.save(imgtmp)
        ret += f"{abspath(imgtmp)}\r\n"
        pd.update_stream(highlight["xref"], img.tobytes())
    # Save the modified PDF to a new file
    pd.save(nfm)
    return ret + f"{abspath(nfm)}\r\n"
# TODO: __main__
def main():
    from sys import argv
    pdf_path = argv[1]
    # coord = eval(argv[2])
    coord = {
        "highlight": [
        {
            "xref": 16,
            "rects": [
                {
                    "x": 0, "y": 0, "h":100, "w":200
                }
            ]
        }
        ]
    }
    print(colorimage(**{
        "pdf_path": pdf_path,
        "highlight": coord["highlight"],
    }))

if __name__ == "__main__":
    main()
